from graphviz import Digraph
from collections import deque
import uuid
from random import sample
class BinaryTree:

    def __init__(self,data=None):
        self.data = data
        self.left = None
        self.right = None
        self.dot = Digraph(comment='Binary Tree')

    def CreateTree(self,data):
        '''
        使用小左大右的顺序建立二叉树
        :param data: 节点数据
        :return:
        '''
        if self.data == None:
            self.data = data
        else:
            if data<self.data:
                if self.left==None:
                    self.left = BinaryTree(data)
                else:
                    self.left.CreateTree(data)
            else:
                if self.right == None:
                    self.right = BinaryTree(data)
                else:
                    self.right.CreateTree(data)

    def preCreateTree(self):
        '''
        需要提前准备一颗空树
        :return:
        '''
        data = int(input("please input the data(-1代表空节点):"))
        if (data == -1):
            return
        else:
            self.data = data
            self.left = BinaryTree()
            self.left.preCreateTree()
            self.right = BinaryTree()
            self.right.preCreateTree()
    def predisplay(self):
        stack = []
        if self.data == None:
            print("空树")
        else:
            stack.append(self)
            while stack:
                cur = stack[-1]
                print(stack.pop().data,end=" ")
                if cur.right!=None:
                    stack.append(cur.right)
                if cur.left!=None:
                    stack.append(cur.left)
    def indisplay(self):
        stack = []
        if self.data == None:
            print("空树")
        else:
            cur = self
            while stack  or cur!=None:
                if cur!=None:
                    stack.append(cur)
                    cur = cur.left
                else:
                    cur = stack.pop()
                    print(cur.data,end=" ")
                    cur = cur.right

    def posdisplay(self):
        #非递归后续遍历
        if self.data == None:
            print("空树")
        else:
            lastVisited = None
            stack = []
            cur = self
            while stack or cur!=None:
                if cur!=None:
                    stack.append(cur)
                    cur = cur.left
                else:
                    peekNode = stack[-1]
                    if peekNode.right!=None and lastVisited != peekNode.right:
                        cur = peekNode.right
                    else:
                        print(peekNode.data,end=" ")
                        lastVisited = stack.pop()

    def leveldisplay(self):
        #层次遍历
        if self.data == None:
            print("空树")
        else:
            cur = self
            lastVisited = deque()
            lastVisited.append(cur)
            while len(lastVisited):
                cur = lastVisited.popleft()
                print(cur.data, end = " ")
                if cur.left!=None:
                    lastVisited.append(cur.left)
                if cur.right!=None:
                    lastVisited.append(cur.right)
    def heightTree(self):
        #计算一棵二叉树的高度
        # 空的树高度为0, 只有root节点的树高度为1
        if self.data is None:
            return 0
        elif self.left is None and self.right is None:
            return 1
        elif self.left is None and self.right is not None:
            return 1 + self.right.heightTree()
        elif self.left is not None and self.right is None:
            return 1 + self.left.heightTree()
        else:
            return 1 + max(self.left.heightTree(), self.right.heightTree())
    def leafnode(self):
        if self.data is None:
            return None
        elif self.left is None and self.right is None:
            print(self.data, end=' ')
        elif self.left is None and self.right is not None:
            self.right.leafnode()
        elif self.right is None and self.left is not None:
            self.left.leafnode()
        else:
            self.left.leafnode()
            self.right.leafnode()
    def VisualTree(self,save_path='./Binary_Tree.gv', label=False):
        # colors for labels of nodes
        colors = ['skyblue', 'tomato', 'orange', 'purple', 'green', 'yellow', 'pink', 'red']

        # 绘制以某个节点为根节点的二叉树
        def print_node(TreeNode, node_tag):
            # 节点颜色
            color = sample(colors, 1)[0]
            if TreeNode.left is not None:
                left_tag = str(uuid.uuid1())  # 左节点的数据
                self.dot.node(left_tag, str(TreeNode.left.data), style='filled', color=color)  # 左节点
                label_string = 'L' if label else ''  # 是否在连接线上写上标签，表明为左子树
                self.dot.edge(node_tag, left_tag, label=label_string)  # 左节点与其父节点的连线
                print_node(TreeNode.left, left_tag)

            if TreeNode.right is not None:
                right_tag = str(uuid.uuid1())
                self.dot.node(right_tag, str(TreeNode.right.data), style='filled', color=color)
                label_string = 'R' if label else ''  # 是否在连接线上写上标签，表明为右子树
                self.dot.edge(node_tag, right_tag, label=label_string)
                print_node(TreeNode.right, right_tag)

        # 如果树非空
        if self.data is not None:
            root_tag = str(uuid.uuid1())  # 根节点标签，唯一标识
            self.dot.node(root_tag, str(self.data), style='filled', color=sample(colors, 1)[0])  # 创建根节点
            print_node(self, root_tag)


        self.dot.view(save_path)  # 保存文件为指定文件,并打开该文件，默认是pdf，可以是png
        # self.dot.render(save_path)# 保存文件为指定文件


if __name__ == '__main__':
    import random
    tree = BinaryTree(5)
    tree.CreateTree(3)
    # tree.CreateTree(2)
    tree.CreateTree(4)
    tree.CreateTree(7)
    tree.CreateTree(8)
    tree.CreateTree(6)
    # tree.predisplay()
    tree.leveldisplay()
    print()
    print(tree.heightTree())
    tree.leafnode()
    tree.VisualTree(label=True)
